import torch
import cv2 as cv
from model.simplenet import SimpleNet
from utils.utils import image_pretreat, image_show


model = SimpleNet()
model.load_state_dict(torch.load('resnet18.pth'))
model.eval()
cap = cv.VideoCapture(0)
cap.set(cv.CAP_PROP_FRAME_WIDTH, 1280)
cap.set(cv.CAP_PROP_FRAME_HEIGHT, 480)
print('Network loading complete.')
while cap.isOpened():
    ret, image = cap.read()
    image = image_pretreat(image[:, :640, :], 224)
    with torch.no_grad():
        image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
        predict = model(image)
    image_show(image[0], predict[0])
    c = cv.waitKey(10)
    if c == ord('q'):
        break
cap.release()
